#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from PIL import Image
import numpy as np

def add_gaussian_noise(image, noise_level=0.03):
    """
    给图像的每个像素添加方差为其值的百分之三的正态分布噪声。

    :param image: 输入图像，PIL.Image 对象
    :param noise_level: 噪声水平，默认为 0.03
    :return: 添加噪声后的图像，PIL.Image 对象
    """
    # 将图像转换为 numpy 数组
    image_array = np.array(image).astype(float)
    
    # 计算每个像素的噪声标准差
    std_dev = image_array * noise_level
    
    # 生成噪声
    noise = np.random.normal(0, std_dev)

    # # 归一化噪声到 [0, 255] 范围
    # noise_normalized = (noise - noise.min()) * 255 / (noise.max() - noise.min())
    # noise_normalized = noise_normalized.astype(np.uint8)
    # # 创建噪声图像并显示
    # noise_img = Image.fromarray(noise_normalized)
    # noise_img.show()
    
    # 添加噪声
    noisy_image_array = image_array + noise
    
    # 确保像素值在有效范围内
    noisy_image_array = np.clip(noisy_image_array, 0, 255).astype(np.uint8)
    
    # 将 numpy 数组转换回 PIL.Image 对象
    noisy_image = Image.fromarray(noisy_image_array)
    
    return noisy_image

#仿真生成的plates图像中部分图像像素不对，通过修改像素值，将前2列和后2列的像素值设置为255，去除孤立的噪点
def refine_plates_image(image):
    # 将图像转换为 numpy 数组
    image_array = np.array(image)
    image_array[:, :2] = 255  # 设置前两列为255
    image_array[:, -2:] = 255  # 设置后两列为255

if __name__ == '__main__':
    source_image_folder = 'D:/project/dataset/dataset_with_noise/TRABECULA_POINTS_Dataset'
    dst_image_folder = 'D:/project/dataset/test/plates'

    files = os.listdir(source_image_folder)
    for file in files:
        if file.endswith('.png'):
            image_path = os.path.join(source_image_folder, file)
            image = Image.open(image_path).convert("L")  # 转换为灰度图像
            noisy_image = add_gaussian_noise(image, noise_level = 0.04)
            dst_imge_path = os.path.join(dst_image_folder, file)
            noisy_image.save(dst_imge_path)
            print(f"添加噪声后的图像已保存到 {dst_imge_path}")